import numpy as np
from sklearn.model_selection import train_test_split

def Count(y,classnum):
    count = np.zeros(classnum,int)
    for i in range(y.size):
        count[y[i]] += 1
    return count

def Gini(y,classnum):
        p = Count(y,classnum)/y.size
        return 1 - p.dot(p)

def GiniIndex(t,x,y,classnum):
        left = x <= t
        right = x > t
        return (left.sum()/x.size)*Gini(y[left],classnum) + (right.sum()/x.size)*Gini(y[right],classnum)

def attributeSelect(X,y,classnum):
        attribute = None
        threshold = None
        gini_index = 1
        for a in range(X.shape[1]):
            x = np.array(X[:,a])
            x.sort()
            for i in range(1,X.shape[0]):
                if x[i-1] != x[i]:
                    t = (x[i-1]+x[i])/2
                    g = GiniIndex(t,X[:,a],y,classnum)
                    if g < gini_index:
                        gini_index = g
                        attribute = a
                        threshold = t
        return attribute,threshold

class Node:
    def __init__(self):
        self.attribute = None
        self.threshold = None
        self.label = None
        self.isLeaf = None
        self.samples = None
        self.left = None
        self.right = None
        return
    
    def decisionMaking(self,x):
        y_p = None
        if self.isLeaf:
            y_p = self.label
        else:
            if x[self.attribute] <= self.threshold:
                y_p = self.left.decisionMaking(x)
            else:
                y_p = self.right.decisionMaking(x)
        return y_p

    def divide(self,X,y,classnum,maxdepth=None,depth=0):
        count = np.zeros(classnum,int)
        for i in range(y.size):
            count[y[i]] += 1
        self.samples = count
        self.isLeaf = True
        self.label = np.argmax(count)
        if (maxdepth == None or depth < maxdepth) and np.count_nonzero(count) > 1:
            a,t = attributeSelect(X,y,classnum)
            if a != None:
                self.attribute = a
                self.threshold = t
                self.isLeaf = False
                self.left = Node()
                self.left.divide(X[X[:,a]<=t,:],y[X[:,a]<=t],classnum,maxdepth,depth+1)
                self.right = Node()
                self.right.divide(X[X[:,a]>t,:],y[X[:,a]>t],classnum,maxdepth,depth+1)
        return

    def postPrun(self,X,y):
        if(y.size == 0):
            precision = 0
        else:
            precision = (y==self.label).sum()
            if(self.left != None):
                left = X[:,self.attribute]<=self.threshold
                right = X[:,self.attribute]>self.threshold
                precision2 = self.left.postPrun(X[left,:],y[left]) + self.right.postPrun(X[right,:],y[right])
                if precision < precision2:
                    precision = precision2
                else:
                    self.isLeaf = True
                    self.left = None
                    self.right = None
        return precision
    
    def prePrunDivide(self,X,y,X_,y_,classnum,maxdepth,depth):
        count = Count(y,classnum)
        self.samples = count
        self.label = np.argmax(count)
        if (maxdepth == None or depth < maxdepth) and np.count_nonzero(count) > 1:
            a,t = attributeSelect(X,y,classnum)
            if a != None:
                left = X[:,a]<=t
                right = X[:,a]>t
                left_ = X_[:,a]<=t
                right_ = X_[:,a]>t
                precision1 = (y_==self.label).sum()
                precision2 = (y_[left_]==np.argmax(Count(y[left],classnum))).sum() + (y_[right_]==np.argmax(Count(y[right],classnum))).sum()
                if precision1 < precision2:
                    self.isLeaf = False
                    self.attribute = a
                    self.threshold = t
                    self.left = Node()
                    self.left.prePrunDivide(X[left,:],y[left],X_[left_,:],y_[left_],classnum,maxdepth,depth+1)
                    self.right = Node()
                    self.right.prePrunDivide(X[right,:],y[right],X_[right_,:],y_[right_],classnum,maxdepth,depth+1)
                else:
                    self.isLeaf = True
            else:
                self.isLeaf = True
        else:
            self.isLeaf = True
        return

    def show(self,fn,tn,depth=0,trace=[]):
        for i in range(depth):
            if trace[i] == True:
                print("|   ",end='')
            else:
                print("    ",end='')
        print("|---",end='')
        print(fn[self.attribute]," <= %.2f"%self.threshold,", samples: ",self.left.samples,end='')
        trace.append(True)
        if(self.left.isLeaf):
            print(",","\033[1;31;36m","label: ",tn[self.left.label],"\033[0m")
        else:
            print('')
            self.left.show(fn,tn,depth+1,trace)
        trace.pop()
        for i in range(depth):
            if trace[i] == True:
                print("|   ",end='')
            else:
                print("    ",end='')
        print("|---",end='')
        print(fn[self.attribute]," > %.2f"%self.threshold,", samples: ",self.right.samples,end='')
        trace.append(False)
        if(self.right.isLeaf):
            print(",","\033[1;31;36m","label: ",tn[self.right.label],"\033[0m")
        else:
            print('')
            self.right.show(fn,tn,depth+1,trace)
        trace.pop()
        return
    

class ClassificationTree:
    def __init__(self,classnum,prun=True,postprun=True,maxdepth=None):
        self.root = Node()
        self.classnum = classnum
        self.prun = prun
        self.postprun = postprun
        self.maxdepth = maxdepth
        return
    
    def fit(self,X,y):
        X = np.array(X)
        y = np.array(y)
        if self.prun:
            X_train,X_test,y_train,y_test = train_test_split(X,y,stratify=y)
            if self.postprun:
                self.root.divide(X_train,y_train,self.classnum,self.maxdepth,0)
                self.root.postPrun(X_test,y_test)
            else:
                self.root.prePrunDivide(X_train,y_train,X_test,y_test,self.classnum,self.maxdepth,0)
        else:
            self.root.divide(X,y,self.classnum,self.maxdepth,0)
        return

    def predict(self,X):
        X = np.array(X)
        y_p = np.empty(X.shape[0],int)
        for i in range(X.shape[0]):
            y_p[i] = (int)(self.root.decisionMaking(X[i,:]))
        return y_p
    
    def show(self,feature_names,target_names):
        print("samples: ",self.root.samples,", labels: ",target_names)
        if(self.root.isLeaf == False):
            self.root.show(feature_names,target_names)
        print('')
        return
